#ifndef OPTIMIZER_H
#define OPTIMIZER_H

#include <map>
#include <string>
#include <vector>
#include <set>
#include <stack>
#include "ASTNode.h"
#include "MatrixOptimizer.h"
#include "OptimizationCostModel.h"
#include "CodeAnalyzer.h"

// 常量折叠优化器
class ConstantFoldingOptimizer {
private:
    bool isConstant(ASTNode* node);
    int evaluateConstant(ASTNode* node);
    bool applyAlgebraicOptimizations(ASTNode* node);

public:
    void optimize(ASTNode* root);
    void optimizeNode(ASTNode* node);
};

// 公共子表达式消除优化器
class CommonSubexpressionOptimizer {
private:
    std::map<std::string, ASTNode*> expressionMap;
    std::map<std::string, std::string> tempVarMap;  // 存储表达式到临时变量的映射
    int tempVarCounter;
    std::string getExpressionKey(ASTNode* node);
    bool isOptimizableExpression(ASTNode* node);
    bool hasVariableReference(ASTNode* node);
    std::string createTempVariable();
    ASTNode* createVarRefNode(const std::string& varName);
    ASTNode* createAssignNode(const std::string& varName, ASTNode* expr);

public:
    CommonSubexpressionOptimizer() : tempVarCounter(0) {}
    void optimize(ASTNode* root);
    void optimizeNode(ASTNode* node);
    void optimizeBlock(ASTNode* blockNode);
    void optimizeFunction(ASTNode* funcNode);
    void optimizeExpressionStatement(ASTNode* exprStmt, std::vector<std::unique_ptr<ASTNode>>& optimizedStatements);
};

// 死代码消除优化器
class DeadCodeEliminator {
private:
    std::set<std::string> usedVariables;  // 记录已使用的变量
    std::set<std::string> definedVariables;  // 记录已定义的变量
    bool isDeadCode(ASTNode* node);
    bool hasSideEffects(ASTNode* node);
    bool isUnreachableCode(ASTNode* node);
    bool isUnusedVariable(ASTNode* node);
    void collectUsedVariables(ASTNode* node);
    void collectDefinedVariables(ASTNode* node);
    void removeDeadCode(ASTNode* parent, ASTNode* deadNode);
    bool isControlFlowStatement(ASTNode* node);
    bool hasReturnStatement(ASTNode* node);

public:
    void optimize(ASTNode* root);
    void optimizeNode(ASTNode* node);
    void optimizeBlock(ASTNode* blockNode);
    void optimizeFunction(ASTNode* funcNode);
};

// 循环优化器
class LoopOptimizer {
private:
    std::map<std::string, ASTNode*> loopInvariants;  // 循环不变量
    std::set<std::string> loopVariables;  // 循环变量
    bool isLoopInvariant(ASTNode* node, const std::set<std::string>& loopVars);
    bool isLoopVariable(const std::string& varName);
    void identifyLoopVariables(ASTNode* loopNode);
    void hoistInvariants(ASTNode* loopNode);
    void strengthReduction(ASTNode* loopNode);
    bool canStrengthReduce(ASTNode* node);

public:
    void optimize(ASTNode* root);
    void optimizeNode(ASTNode* node);
    void optimizeLoop(ASTNode* loopNode);
};

// 寄存器分配优化器
class RegisterAllocationOptimizer {
private:
    std::map<std::string, int> variableRegisters;  // 变量到寄存器的映射
    std::map<int, std::string> registerVariables;  // 寄存器到变量的映射
    std::stack<int> freeRegisters;  // 空闲寄存器栈
    int nextRegister;
    int maxRegisters;

    void initializeRegisters();
    int allocateRegister(const std::string& varName);
    void freeRegister(int reg);
    void spillRegister(int reg);
    void analyzeLiveness(ASTNode* node);
    std::set<std::string> getLiveVariables(ASTNode* node);
    void optimizeRegisterUsage(ASTNode* node);

public:
    RegisterAllocationOptimizer() : nextRegister(0), maxRegisters(32) { initializeRegisters(); }
    void optimize(ASTNode* root);
    void optimizeNode(ASTNode* node);
    void optimizeFunction(ASTNode* funcNode);
};

// 内联优化器
class InliningOptimizer {
private:
    std::map<std::string, ASTNode*> functionDefinitions;  // 函数定义
    std::map<std::string, int> functionSizes;  // 函数大小
    int maxInlineSize;
    int inlineDepth;

    void collectFunctionDefinitions(ASTNode* root);
    int calculateFunctionSize(ASTNode* funcNode);
    bool shouldInline(const std::string& funcName);
    ASTNode* createInlinedFunction(ASTNode* callNode, ASTNode* funcDef);
    void substituteParameters(ASTNode* inlinedBody, const std::vector<ASTNode*>& args);
    std::string generateUniqueName(const std::string& base);

public:
    InliningOptimizer() : maxInlineSize(50), inlineDepth(0) {}
    void optimize(ASTNode* root);
    void optimizeNode(ASTNode* node);
    void optimizeFunction(ASTNode* funcNode);
};

// 向量化优化器
class VectorizationOptimizer {
private:
    std::vector<std::string> vectorizableOperations;  // 可向量化的操作
    bool isVectorizableLoop(ASTNode* loopNode);
    bool isVectorizableOperation(ASTNode* node);
    void identifyVectorizablePatterns(ASTNode* node);
    ASTNode* createVectorizedLoop(ASTNode* originalLoop);
    bool hasDataDependencies(ASTNode* loopBody);
    int getLoopIterationCount(ASTNode* loopNode);

public:
    VectorizationOptimizer();
    void optimize(ASTNode* root);
    void optimizeNode(ASTNode* node);
    void optimizeLoop(ASTNode* loopNode);
};

// 主优化器
class Optimizer {
private:
    ConstantFoldingOptimizer constantFolder;
    CommonSubexpressionOptimizer cseOptimizer;
    DeadCodeEliminator deadCodeEliminator;
    MatrixOptimizer matrixOptimizer;
    LoopOptimizer loopOptimizer;
    RegisterAllocationOptimizer registerAllocator;
    InliningOptimizer inliningOptimizer;
    VectorizationOptimizer vectorizationOptimizer;
    OptimizationCoordinator coordinator;
    CodeAnalyzer codeAnalyzer;
    bool useCoordinatedOptimization;
    bool useAdaptiveOptimization;

public:
    Optimizer() : useCoordinatedOptimization(false), useAdaptiveOptimization(true) {}

    void optimize(ASTNode* root);
    void optimizeConstantFolding(ASTNode* root);
    void optimizeCommonSubexpressions(ASTNode* root);
    void optimizeDeadCodeElimination(ASTNode* root);
    void optimizeMatrixOperations(ASTNode* root);
    void optimizeLoopOperations(ASTNode* root);
    void optimizeRegisterAllocation(ASTNode* root);
    void optimizeInlining(ASTNode* root);
    void optimizeVectorization(ASTNode* root);

    // 协调优化相关方法
    void enableCoordinatedOptimization(bool enable);
    void setOptimizationStrategy(const std::vector<OptimizationType>& strategy);
    void setOptimizationLevel(int level);

    // 自适应优化相关方法
    void enableAdaptiveOptimization(bool enable);
    void analyzeAndOptimize(ASTNode* root);
    OptimizationStrategy getOptimizationStrategy(ASTNode* root);
    void applyAdaptiveOptimizations(ASTNode* root, const OptimizationStrategy& strategy);
    void applyOptimization(ASTNode* root, OptimizationType optimization);
};

#endif // OPTIMIZER_H
